#!/usr/bin/env python
# coding=utf-8
# CopyRight(c) 2019 Inspur Electronic Information Industry Co.,Ltd

"""
    secchecksuite/main.py
    :CopyRight(c) 2019 Inspur Electronic Information Industry Co.,Ltd
"""
import paramiko
import argparse
import atexit
import datetime
import os
import sys
import tempfile

import constant as const
import seccheck_log
import export
import plugin
import sqlitedb
import aes_cbc
import utils
import check
import config

# global var
# sqlite database connetion
conn = None
# AES key
AES_KEY = ''
# secchecksuite config info
config_info = dict()
PythonVersionObj = object()


def add_argument():
    """
    Parse and return the arguments.
    Returns:
        args: return the argument_list
    Raises:
        IOError: Input or output error
    """
    parser = argparse.ArgumentParser(
        description='security scan and reinforce CentOS 7.X')
    parser.add_argument(
        '-v', '--version',
        action='version',
        version='sec_check_suite v{}'.format(const.VERSION)
    )
    parser.add_argument(
        '--plugins_config', type=str, help='plugins_config file of secchecksuite',
        default='./../conf/secchecksuite.json'
    )
    parser.add_argument(
        '--cleanup', action='store_true', help='clean database tables'
    )
    group = parser.add_mutually_exclusive_group()
    group.add_argument(
        '--encrypt', type=str, help='specify plaintext to be encrypted'
    )
    group.add_argument(
        '--decrypt', type=str, help='specify ciphertext to be decrypted'
    )
    group.add_argument(
        '--rollback', help='rollback all reinforcement items performed',
        action='store_true'
    )
    group.add_argument(
        '--generate', action='store_true', help='generate report'
    )
    parser.add_argument(
        '--key', type=str, help='specify key for encryption/decryption,'
                                'at lest 8 digits')
    parser.add_argument(
        '--daemon', help='Start the daemon process ', action='store_true'
    )
    parser.add_argument('-c', '--config', help='config configure file information', action='store_true')
    parser.add_argument('-m', '--module', help='plugin module to execute,\n '
                                               'Input Format:[entryPoints, module],'
                                               'entryPoints match the key of entrypoints in secchecksuite.json,'
                                               'module match the id in plugin json.  \n'
                                               'entryPoints: chose [all, docker, host, k8s, third].\n'
                                               'module host: chose [all, access_control, anti_malware, authentication, intrusion_guard, security_audit, system];\n'
                                               'module docker chose [all, docker_daemon, docker_dockered, docker_file, docker_host, docker_operate, docker_swarm];\n'
                                               'module k8s chose [all, apiserver, controller, etcd, kubelet, master, node, scheduler]\n'
                                               'module third: [docker_bench_security]\n',
                        default='ALL')
    args = vars(parser.parse_args())
    return args


def init_args():
    """
    add key into args.
    Returns:
        args: return the argument_list
    Raises:
        IOError: Input or output error
    """
    try:
        args = add_argument()
        # 默认配置文件路径
        if not os.path.isfile(args.get('plugins_config', None)):
            error_des = '%s %s %s' % ('config file ', args.get('plugins_config', None), 'not found')
            raise IOError(error_des)
        if not args.get('generate', None) and \
                not args.get('cleanup', None) and \
                not check.validate_input_key(args.get('key', None)):
            args['key'] = aes_cbc.get_input_secret_key(PythonVersionObj, args.get('key', None))
        global AES_KEY
        AES_KEY = args.get('key', None)
        return args
    except (SystemExit, BaseException) as err:
        if check.check_normal_exit(err):
            utils.program_exit(sys.argv[0], 0)
        print("parse args error:{}".format(err))
        utils.program_exit(sys.argv[0], -1)


def init_sqlite():
    """
    init sqlite database
    Returns:
        conn:sqlite database connetions
    Raises:
        BaseException:  An error occurred createing a table
    """
    try:
        seccheck_log.sec_check_logger.info('Init sqlite database begin!')
        # 获取数据库文件路径
        db_file_config = config_info.get('dbfilepath', None)
        if db_file_config is not None and db_file_config != '':
            db_file_path = os.path.join(db_file_config, const.SQLITE_DB_NAME)
        else:
            db_file_path = const.SQLITE_DB_NAME
        # 连接数据库
        global conn
        conn = sqlitedb.connect_sqlite(db_file_path)
        sqlitedb.set_sqlite_table_info(conn)
    except BaseException as e:
        seccheck_log.sec_check_logger.error('init_sqlite occur some error {}'.format(e))
        utils.program_exit(sys.argv[0], -1)


def daemon_init(pid_file=None):
    """
    创建守护进程
    Args:
        param pid_file: 保存进程id的文件
    Returns:
        None
    """
    # 从父进程fork一个子进程出来
    pid = os.fork()
    # 子进程的pid一定为0，父进程大于0
    if pid:
        # 退出父进程，sys.exit()方法比os._exit()方法会多执行一些刷新缓冲工作
        sys.exit(0)
    # 子进程默认继承父进程的工作目录，最好是变更到根目录，否则回影响文件系统的卸载
    # os.chdir('./')
    # 子进程默认继承父进程的umask（文件权限掩码），重设为0（完全控制），以免影响程序读写文件
    os.umask(0)
    # 让子进程成为新的会话组长和进程组长
    os.setsid()
    # 注意了，这里是第2次fork，也就是子进程的子进程，我们把它叫为孙子进程
    _pid = os.fork()
    if _pid:
        # 退出子进程
        sys.exit(0)
    # 此时，孙子进程已经是守护进程了，接下来重定向标准输入、输出、错误的描述符(是重定向而不是关闭, 这样可以避免程序在 print 的时候出错)
    # 刷新缓冲区先，小心使得万年船
    sys.stdout.flush()
    sys.stderr.flush()

    # dup2函数原子化地关闭和复制文件描述符，重定向到/dev/nul，即丢弃所有输入输出
    with open('/dev/null') as read_null, open('/dev/null', 'w') as write_null:
        os.dup2(read_null.fileno(), sys.stdin.fileno())
        os.dup2(write_null.fileno(), sys.stdout.fileno())
        os.dup2(write_null.fileno(), sys.stderr.fileno())

    # 写入pid文件
    if pid_file:
        with open(pid_file, 'w+') as f:
            f.write(str(os.getpid()))
        # 注册退出函数，进程异常退出时移除pid文件
        atexit.register(os.remove, pid_file)


def application_instance(file):
    """
    run only one instance
    Returns:

    """
    basename = os.path.splitext(os.path.abspath(file))[0].replace(
        "/", "").replace(":", "").replace("\\", "-") + '-%s' % 'security' + '.lock'
    lockfile = os.path.normpath(tempfile.gettempdir() + '/' + basename)
    try:
        if os.path.exists(lockfile):
            print("Another instance is already running, quitting.")
            sys.exit(0)
        else:
            open(lockfile, 'w')
    except BaseException as e:
        sys.exit(0)


def main():
    """
    security scan/ defend
    Raises:
        IOError: Input or Output error
        json.decoder.JSONDecodeError: JSON decode error
        KeyError: key board error
        BaseException: An error when init intermediate file
    """
    # begin initialization
    time_begin = datetime.datetime.now()

    global PythonVersionObj
    PythonVersionObj = utils.judge_python()
    PythonVersionObj.set_default_encoding()

    # 信号管理
    check.check_signal()

    # 单实例处理
    application_instance(sys.argv[0])

    # 获取命令行参数
    args = init_args()

    # 处理字符串加解密
    aes_cbc.aes_enc_decrypt(args, AES_KEY)

    # 守护进程处理
    if args.get('daemon', None):
        daemon_init()

    # 加载日志/插件集配置文件
    try:
        global config_info
        config_info = utils.load_config_file(args.get('plugins_config', None))
        if not check.check_not_none_dict(config_info):
            utils.program_exit(sys.argv[0], -1)
        check.check_config_info(config_info)
    except BaseException as err:
        print("load_main_config error:{}".format(err))
        utils.program_exit(sys.argv[0], -1)

    # 配置脚本文件
    if args.get('config', None):
        config_obj = config.SecurityConfig(PythonVersionObj, AES_KEY, config_info)
        config_obj.config()

    # 日志初始化
    tmp_log_file = seccheck_log.init_logger(config_info)

    seccheck_log.sec_check_logger.info('start load main config')

    # 产生扫描报告
    if args.get('generate', None):
        export.export_result(tmp_log_file, file_type=config_info.get('exporttype', 'json'))
        utils.program_exit(sys.argv[0], 0)

    # 数据库初始化
    init_sqlite()

    # 清理数据库中信息
    if args.get('cleanup', None):
        sqlitedb.clean_database(conn, config_info)

    # 中间临时文件初始化
    seccheck_log.init_intermediate_file(config_info.get('loginfo', None))

    # 执行插件
    run_plugin(args, time_begin)

    # 结果导出
    export.export_result(tmp_log_file, file_type=config_info.get('exporttype', None))

    # end
    time_end = datetime.datetime.now()
    time = time_end - time_begin
    seccheck_log.sec_check_logger.info('exec all plugins finish. cost {}s'.format(time.total_seconds()))
    utils.program_exit(sys.argv[0], 0)


def run_plugin(args, time_begin):
    # 执行/回退 插件内容
    if args.get('rollback', None):
        if not check.check_database_valid(config_info, const.SQLITE_DB_NAME):
            seccheck_log.sec_check_logger.error(
                'Database is not exist, cannot execute rollback. Please execute scan or defend first.')
            utils.program_exit(sys.argv[0], 0)
        plugin_lists = sqlitedb.getinfo_from_database(conn)
        if not plugin_lists or not len(plugin_lists):
            seccheck_log.sec_check_logger.error(
                'Database is null, cannot execute rollback. Please execute scan or defend first.')
            utils.program_exit(sys.argv[0], 0)
        plugin.rollback_all(config_info, conn, AES_KEY)
    else:
        module = init_module(args)
        # parse and save plugins
        plugin.parse_plugins(module, config_info, conn, AES_KEY)
        # initialization finish
        time_init = datetime.datetime.now()
        time = time_init - time_begin
        seccheck_log.sec_check_logger.info('success initialization. cost {}s'.format(time.total_seconds()))

        # 执行插件
        plugin.exec_plugins(config_info, conn, AES_KEY)


def init_module(args):
    args_module = args.get('module', None)
    config_obj = config.SecurityConfig(PythonVersionObj, AES_KEY, config_info)
    module, flag = config_obj.complete_valid_module(args_module)
    if not flag:
        module = config_obj.get_module(False)
    if not module:
        module = dict()
        module['entry'] = 'all'
        module['module'] = 'all'
        return module
    if not module.get('entry', None):
        module['entry'] = 'all'
    if not module.get('module', None):
        module['module'] = 'all'
    return module


if __name__ == "__main__":
    try:
        main()
    except (SystemExit, BaseException) as e:
        if check.check_normal_exit(e):
            utils.program_exit(sys.argv[0], 0)
        print('main exit,detail error:{}.'.format(e))
        utils.program_exit(sys.argv[0], -1)
